import re
from model.resnet import ModifiedResNet
import torch
from torch import Tensor
from PIL import Image
from PIL.Image import NONE
from model.data_helper import ImageFactory
from model.trainer import Trainer, tensor_to_PIL
from model import trainer
import numpy as np
import config as configpy
from abc import abstractmethod
import os
import config
from numpy.core.numerictypes import issubdtype

TERMINATE = -1
DUPLICATE_LAST = -2

CLS_TO_ACTION_DICT = {
    "n": 0,
    "w": 2,
    "a": 4,
    "s": 5,
    "d": 3,
    "f": 1,
    "e": 11,
    "q": 12,
    "t": TERMINATE,
    "p": DUPLICATE_LAST
}

ACTION_TO_CLS_DICT = {
    0: "None",
    2: 'Up',
    4: 'Left',
    5: "Down",
    3: "Right",
    1: "Jump",
    11: "RightJump",
    12: "LeftJump",
    TERMINATE: "TERMINATE"
}

ACTION_TO_CLS_DICT = {
    0: "None",
    2: 'Up',
    4: 'Left',
    5: "Down",
    3: "Right",
    1: "Jump",
    11: "RightJump",
    12: "LeftJump",
    TERMINATE: "TERMINATE"
}

IDX_TO_CLS_DICT = {
    0: '0',
    1: '1',
    2: '2',
    3: '3',
    4: '4',
    5: '5',
    6: '11',
    7: '12',
}

NUM_ACTION = len(CLS_TO_ACTION_DICT) - 1

device = "cuda" if torch.cuda.is_available() else "cpu"


def file_append(file_path, line):
    flie_folder = os.path.dirname(os.path.abspath(file_path))
    if not os.path.exists(flie_folder):
        os.makedirs(flie_folder)

    with open(file_path, 'a+') as f:
        f.write(line)


def label_file_append(img_path, label):
    file_append(config.LABEL_PATH, str(img_path) + str(label))


def get_action():
    has_command = False
    times = 1
    r = "^(?P<times>[0-9]*)(?P<action>[a-z]{1})$"
    while has_command is False:
        # action = int(input('input action '))
        cmd_s = input("input action").strip()
        cmd_group = re.search(r, cmd_s)
        if cmd_group == None:
            print("Input NUM+ACTION")
            continue
        else:
            cmd_action = cmd_group.group('action')
            cmd_times = cmd_group.group('times')

        if cmd_times == '':
            times = 1
        else:
            times = int(cmd_times)
        action = CLS_TO_ACTION_DICT.get(cmd_action)
        has_command = True
        if action is None:
            print(
                'num + input n -> none , w,a,s,d -> directions, f -> fire, q,e -> left/right fire; t -> TERMINATE'
            )
            has_command = False
    return action, times


class DaggerAgent:
    def __init__(self, ):
        pass

    @abstractmethod
    def select_action(self, ob):
        pass


# TODO MyAgent


# here is an example of creating your own Agent
class MyAgent(DaggerAgent):
    def __init__(self, model, cls_txt_path):
        super(DaggerAgent, self).__init__()
        # init your model
        self.model = model
        self.image_factory = ImageFactory(cls_txt_path)

    # train your model with labeled data
    def update(self,
               imgs_labels_txt_filepath,
               trainer: Trainer,
               increment_trainning=True):

        if increment_trainning == False:
            # ReTrainning
            self.model = ModifiedResNet()

        train_dataloader, val_dataloader, test_dataloader, idx_to_cls_dict = self.image_factory.get_iterable_image_dataloaders(
            imgs_labels_txt_filepath, batch_size=trainer.configs.BATCH_SIZE)
        self.idx_to_cls_dict = idx_to_cls_dict
        trainer.train(self.model, train_dataloader, val_dataloader,
                      test_dataloader, configpy.TIME + configpy.TASK_SUFFIX)
        num_examples = len(train_dataloader.dataset)
        return num_examples

    # select actions by your model
    def select_action(self, obs, enable_print=True):
        self.model.eval()
        img_array = Image.fromarray(obs)
        img_tensor = ImageFactory.image_to_tensor(img_array=img_array)
        # save the obs
        _img_for_infer = tensor_to_PIL(img_tensor)
        _img_for_infer.save("./check_infer.jpg")

        img_tensor = Tensor(img_tensor).unsqueeze(0)

        out_logits = self.model(img_tensor)
        output_argmaxed = out_logits.argmax(dim=1)
        action = self.idx_to_cls_dict.get(output_argmaxed.item())
        if enable_print == True:
            print("原始输出为 {0:d}".format(output_argmaxed.item()))
            print(self.idx_to_cls_dict)
            print("\t对应 {0}".format(action))
        return int(action)

    # def eval(self,img_labels_txt_path):
    #     self.model.eval()
    #     test_dataloader, _,_,_ = ImageFactory.get_iterable_image_dataloaders(imgs_and_labels_txt_path=img_labels_txt_path,batch_size= 1)
    #     for batch_idx, example in enumerate(test_dataloader):
    #         inputs, labels = example

    #         # forward + backward + optimize
    #         outputs = model(inputs)
    #         labels = labels.to(device)

    #         loss = loss_function(outputs, labels)
    #         loss.backward()
    #         optimizer.step()
    #         output_argmaxed = outputs.argmax(dim=1)

    #         # Test
    #         # raise NotImplementedError("查看数据集对不对")
    #         if enable_test == True:
    #             test_idx = 1
    #             test_img_tensor = inputs[test_idx]
    #             test_img = tensor_to_PIL(test_img_tensor)
    #             utils.path_check("./dagger_outputs/")
    #             test_img.save(
    #                 './dagger_outputs/output_{0}_label_{1}.jpg'.format(
    #                     output_argmaxed[test_idx].item(),
    #                     labels[test_idx].item()))

    #         correct = (output_argmaxed == labels).float().sum()
    #         acc = correct / inputs.size()[0]
    #         batch_acces.append(acc.item())
    #         batch_losses.append(loss.item())

    #         output_classes = [
    #             int(Dagger.IDX_TO_CLS_DICT.get(_.item()))
    #             for _ in output_argmaxed
    #         ]
    #         output_actions = [
    #             Dagger.ACTION_TO_CLS_DICT.get(_) for _ in output_classes
    #         ]
    #         labels_classes = [
    #             int(Dagger.IDX_TO_CLS_DICT.get(_.item())) for _ in labels
    #         ]
    #         labels_actions = [
    #             Dagger.ACTION_TO_CLS_DICT.get(_) for _ in labels_classes
    #         ]

    #         print("output", output_actions)
    #         print("label", labels_actions)

    #         print("iter : {0}, loss : {1:.2f}, acc : {2:.2f}".format(
    #             batch_idx, loss.item(), acc.item()))

    #     return np.mean(batch_acces), np.mean(batch_losses)


# here is an example of creating your own Agent
class ExampleAgent(DaggerAgent):
    def __init__(self, necessary_parameters=None):
        super(DaggerAgent, self).__init__()
        # init your model
        self.model = None

    # train your model with labeled data
    def update(self, data_batch, label_batch):
        self.model.train(data_batch, label_batch)

    # select actions by your model
    def select_action(self, data_batch):
        # FIXME Predict
        label_predict = CLS_TO_ACTION_DICT.get('r')
        # label_predict = self.model.predict(data_batch)
        return label_predict
